package com.iot.ig.websocket.handler;

import com.iot.ig.cache.LocalUserCache;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.*;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @author: jingwei
 * @description: WebSocketHandler
 */
@Slf4j
@Component
public class IotWebSocketHandler extends TextWebSocketHandler {

    private static final String WS_SESSION_KEY = "username";

    private final Map<String, WebSocketSession> userSessions = new HashMap<>();

    private final CopyOnWriteArraySet<WebSocketSession> sessions = new CopyOnWriteArraySet<>();

    /**
     * 连接建立成功，可以进行消息传输，可以通过session.sendMessage()发送消息
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String username = (String) session.getAttributes().get(WS_SESSION_KEY);
        userSessions.put(username, session);
        sessions.add(session);
        // 模拟用户缓存，实际项目中可删除下面一行代码
        mockUserCache(username);
        log.info("建立websocket连接: {}", session.getId());
    }

    private void mockUserCache(String username) {
        List<String> cacheList = LocalUserCache.getCache(WS_SESSION_KEY);
        if(cacheList == null) {
            cacheList = new ArrayList<>();
        }
        cacheList.add(username);
        LocalUserCache.addCache(WS_SESSION_KEY, cacheList);
    }

    /**
     * 接收到消息
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String username = (String) session.getAttributes().get(WS_SESSION_KEY);
        log.info("收到websocket消息: username={}，payload={}", username, message.getPayload());
        session.sendMessage(new TextMessage("Echo: " + message.getPayload()));
    }

    /**
     * 连接关闭时候调用
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String username = (String) session.getAttributes().get(WS_SESSION_KEY);
        userSessions.remove(username);
        sessions.remove(session);
        log.info("websocket连接关闭: sessionId={}, username={} ", session.getId(), username);
    }

    /**
     * 广播消息
     * @param message
     */
    public void broadcastMessage(String message) {
        for (WebSocketSession session : sessions) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(message));
                }
            } catch (Exception e) {
                log.error("Error sending message to session: {}", session.getId());
                e.printStackTrace();
            }
        }
    }

    /**
     * 给某个用户发送消息
     * @param username
     * @param message
     */
    public void sendMessageToUser(String username, String message) {
        log.info("websocket send message to user: {}", username);
        WebSocketSession session = userSessions.get(username);
        log.info("session: {}", session);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message));
                log.info("websocket send message success!");
            } catch (Exception e) {
                log.error("Error sending message to user: {}", username);
                e.printStackTrace();
            }
        } else {
            log.error("User not found or session is not open: {}", username);
        }
    }

    /**
     * 批量发送消息
     *
     * @param usernames
     * @param message
     */
    public void sendMessageToUsers(String[] usernames, String message) {
        log.info("批量发送消息，usernames: {}", Arrays.toString(usernames));
        // 判断usernames是否为空
        if (usernames == null || usernames.length == 0) {
            log.info("username is empty!");
            return;
        }
        for (String username : usernames) {
            sendMessageToUser(username, message);
        }
    }
}
